"""
https://pytorch.org/vision/stable/datasets.html#torchvision.datasets.DatasetFolder

FAILED AND NO ADEQUATE TIME TO STUDY IT.
"""
import torch as pt
import torchvision as ptv
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import os
import sys
import re
import numpy as np

pt.manual_seed(777)

BATCH_SIZE = 16

data_dir = '../../../../large_data/DL2/_many_files/cifar2_fast/train'
auto_re = re.compile(r'^.*[\\/]automobile[\\/][^\\/]+$')


def load_img(path):
    data = plt.imread(path)
    data = data.transpose(0, 2)
    if auto_re.match(path):
        y = 0
    else:
        y = 1
    y = pt.tensor(y)
    return data, y



def is_valid_file(path):
    ext = os.path.splitext(path)[1].lower()
    if ext == '.jpg':
        return True
    else:
        return False


train_ds = ptv.datasets.DatasetFolder(root=data_dir, loader=load_img, is_valid_file=is_valid_file, transform=ptv.transforms.ToTensor())
train_dl = DataLoader(dataset=train_ds, batch_size=BATCH_SIZE, shuffle=True)

plt.figure(figsize=[12, 6])
spr = 4
spc = 8
spn = 0


def show_pics(dl):
    global spn
    for bx, by in dl:
        for i, bxi in enumerate(bx):
            print(bxi.shape)
            sys.exit(0)
            spn += 1
            if (spn > spr * spc):
                break
            plt.subplot(spr, spc, spn)
            bxi = bxi.transpose(0, 1)
            plt.imshow(bxi)
            plt.axis('off')
            cls_id = by[i].item()
            plt.title(str(cls_id))
        if (spn > spr * spc):
            break


show_pics(train_dl)
